# -*- coding: utf-8 -*-
from jili.data.db import getk1mbar_pkl,getkline1d
from jili.price.kline import klinenm
from jili.price.ta import ta as tatoolbox
import pandas as pd
import talib
from talib import abstract
def getbasebar(obj):
    return (getk1mbar_pkl(obj),getkline1d(obj,"list"))
def callta_df(kbard,ta_arg):
    for i in ta_arg:
        cmd=i["cmd"]
        call=abstract.Function(cmd)
        if "arg" in i.keys():
            call.set_parameters(i["arg"])
        try:
            if call.set_input_arrays(kbard):
                if "out" in i.keys():
                    o=i["out"]
                    n=len(o)
                    if n==1:
                        k=o[0]
                        if k!="real":
                            k=cmd+"."+k
                        else:
                            k=cmd
                        kbard[k]=call.outputs
                    else:
                        for j in range(n):
                            k=o[j]
                            if k!="real":
                                k=cmd+"."+o[j]
                            else:
                                k=cmd
                            kbard[k]=call.outputs[call.output_names[j]]
                else:
                    if len(call.output_names)==1:
                        kbard[cmd]=call.outputs
                    else:
                        name=[]
                        for j in call.output_names:
                            if j=="real":
                                name.append(cmd)
                            else:
                                name.append(cmd+"."+j)
                        tmp=call.outputs
                        tmp.columns=name
                        kbard.merge(tmp,left_index=True, right_index=True, how='left')
        except Exception as e:
            print(call.info)
            print(e)
    return kbard
def callta(kbar,ta_arg):
    if isinstance(kbar,list):
        kbard=pd.DataFrame(kbar)
        kbard.index=kbard["timekey"]
        kbard=callta_df(kbard,ta_arg)
        return kbard.to_dict("records")
    else:
        return callta_df(kbar,ta_arg)
def getklineset_df(obj,ta_args=[],timeperiod=[5,15]):
    bar_1m,bar_1d=getbasebar(obj)
    bars={}
    rst=[]
    tas={}
    if ta_args:
        for i in ta_args:
            tas["1m"]=[]
            tas["1m"].append(tatoolbox(i))
    for i in timeperiod:
        k=str(i)+"m"
        bars[k]=klinenm(obj,period=i)
        if ta_args:
            for i in ta_args:
                tas[k] = []
                tas[k] .append(tatoolbox(i))
    for i in bar_1m:
        t={}
        t["timekey"]=i["timekey"]
        for k ,v in i.items():
            if k !="timekey":
                k0=k+".1m"
                t[k0]=v
        if tas:
            for ta0 in tas["1m"]:
                tarst=ta0.onbar(i)
                for k,v in tarst.items():
                    k0=k+".1m"
                    t[k0]=v
        for k,bar in bars.items():
            b=bar.onbar(i)
            for k1, v1 in b.items():
                k0 = k1 + "."+k
                t[k0] = v1

            if tas:
                for ta0 in tas[k]:
                    tarst = ta0.onbar(b)
                    for k1, v1 in tarst.items():
                        k0 = k1 + "."+k
                        t[k0] = v1
        rst.append(t)
    rst=pd.DataFrame(rst)
    rst.index=rst["timekey"]
    return rst
def getklineset(obj,ta_args=[],timeperiod=[5,15]):
    bar_1m,bar_1d=getbasebar(obj)
    bars={}
    rst=[]
    tas={}
    if ta_args:
        for i in ta_args:
            tas["1m"]=[]
            tas["1m"].append(tatoolbox(i))
    for i in timeperiod:
        k=str(i)+"m"
        bars[k]=klinenm(obj,period=i)
        if ta_args:
            for i in ta_args:
                tas[k] = []
                tas[k] .append(tatoolbox(i))
    for i in bar_1m:
        set={}
        set["timekey"]=i["timekey"]
        t={}
        t["timekey"]=i["timekey"]
        for k ,v in i.items():
            t[k]=v
        if tas:
            for ta0 in tas["1m"]:
                tarst=ta0.onbar(i)
                for k,v in tarst.items():
                    t[k]=v
        set["1m"]=t
        for k,bar in bars.items():
            t={}
            b=bar.onbar(i)
            for k1, v1 in b.items():
                t[k1] = v1
            if tas:
                for ta0 in tas[k]:
                    tarst = ta0.onbar(b)
                    for k1, v1 in tarst.items():
                        t[k1] = v1
            set[k]=t
    rst.append(set)
    return rst
#性能测试
if __name__ == "__main__":
    ta_args = [{"cmd": "BBANDS", "out": ["upperband", "middleband", "lowerband"],
                "arg": {"timeperiod": 9, "nbdevup": 3, "nbdevdn": 3, "matype": 0}, "batch": 9}]
    a = getklineset("rb1805", ta_args, [3, 5, 15, 30])
        
"""
#性能测试
if __name__ == "__main__":
    from jili.data.db import getk1m_pkl
    from jili.price.tas import tas as tatoolbox
    import talib
    import time
    from talib import abstract
    obj="rb1805"
    k1m=getk1m_pkl(obj)
    start=time.perf_counter()
    k1md=pd.DataFrame(k1m)
    k1md.index=k1md["timekey"]
    k1md1=callta(k1md,[{"cmd": "BBANDS", "out": ["upperband", "middleband", "lowerband"], "arg": {"timeperiod": 9, "nbdevup": 3, "nbdevdn": 3, "matype": 0}, "batch": 5}])
    end=time.perf_counter()
    print("case 1:",end-start)
    start=time.perf_counter()

    tl=tatoolbox({"cmd": "BBANDS", "out": ["upperband", "middleband", "lowerband"], "arg": {"timeperiod": 9, "nbdevup": 3, "nbdevdn": 3, "matype": 0}, "batch": 5})
    #a=k1m[:5]
    #for i in a:
    for i in k1m:
        r=tl.onbarcalc(i)
        #print(r)
        for k,v in r.items():
            i[k]=v
    end=time.perf_counter()
    print("case 2:",end-start)
if __name__=="__main__":
    import talib
    import json
    k1m=getk1m_pkl("rb1805")
    k1md=pd.DataFrame(k1m)
    k1md.index=k1md["timekey"]
    args={}
    arg1=[]
    for i in talib.get_functions():
        t={}
        t["cmd"]=i
        call=abstract.Function(i)
        t["out"]=call.output_names
        tt={}
        for k,v in call.get_parameters().items():
            tt[k]=v
            t["arg"]=tt
        args[i]=t
        arg1.append(t)
    k1md1=callta(k1md,arg1)
    aa={}
    for i in k1md.columns:
        cmd = i.split(".")[0]
        if cmd not in aa.keys():
            aa[cmd]=k1md[i].isna().sum()
        else:
            n=k1md[i].isna().sum()
            if n>aa[cmd]:
                aa[cmd]=n
    f=open("ta_args.json","w")
    for k ,v in aa.items():
        if k in args.keys():
            tas=args[k]
            batch=int(v)
            if batch>1000:
                tas["batch"]=1
            else:
                tas["batch"]=int(v)+1
    json.dump(list(args.values()),f)
    f.close()
    f=open("ta_args.txt","w")
    for k,v in talib.get_function_groups().items():
        f.writelines("-------------------------"+k + "---------------------\n")
        for i in v:
            f.writelines(json.dumps(args[i])+"\n")
    f.close()
    #return args,arg1
"""